package com.lwc.backendsinglestructure.llm.vector;

import com.lwc.backendsinglestructure.llm.constant.MilvusConstants;
import com.lwc.backendsinglestructure.llm.enums.EmbeddingModelEnum;
import io.milvus.param.R;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;

import java.util.List;

public class ArticleMilvusService<T> extends AbstractMilvusService<T>{

    @Override
    public boolean writeCollection(T data) {
        return false;
    }

    /**
     * 向量维度
     * @return
     */
    @Override
    protected int dimension() {
        return EmbeddingModelEnum.ALIBABA.getDimension();
    }

    /**
     * 搜索
     * @param query 用户查询
     * @param topK
     * @param gtScore
     * @return
     */
    @Override
    public List<R> search(String query, int topK, float gtScore) {
        // 对用户query做embedding：这里embedding模型需要同同原数据embed采用同模型
        float[] queryVector = super.embed(query);
        return List.of();
    }

    /**
     * 定义向量数据库索引
     * @return
     */
    @Override
    protected List<IndexParam> indexParams() {
        IndexParam summaryVectorIndex = IndexParam.builder()
                .fieldName("summary_vector")
                // 余弦相似度
                .metricType(IndexParam.MetricType.COSINE)
                .build();
        IndexParam contentVectorIndex = IndexParam.builder()
                .fieldName("content_vector")
                // 余弦相似度
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        return List.of(summaryVectorIndex, contentVectorIndex);
    }

    @Override
    protected String collcetionName() {
        return MilvusConstants.CollectionName.COLLECTION_ARTICLE.name();
    }

    /**
     * 定义业务字段 + 向量字段
     * @return
     */
    @Override
    protected CreateCollectionReq.CollectionSchema schema() {
        CreateCollectionReq.CollectionSchema schema = clientV2.createSchema();
        // 主键
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.VarChar)
                .isPrimaryKey(true)
                .autoID(true)
                .build());
        /**
         * 业务字段
         */
        schema.addField(AddFieldReq.builder()
                .fieldName("pk")  // 文章主键
                .dataType(DataType.VarChar)
                .maxLength(64)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("source_type")
                .dataType(DataType.VarChar)
                .maxLength(64)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("summary")
                .dataType(DataType.VarChar)
                .maxLength(500)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("content")
                .dataType(DataType.VarChar)
                .maxLength(65535)
                .build());
        /**
         * 向量字段
         */
        schema.addField(AddFieldReq.builder()
                .fieldName("summary_vector")
                .dataType(DataType.FloatVector)
                .dimension(this.dimension())
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("content_vector")
                .dataType(DataType.FloatVector)
                .dimension(this.dimension())
                .build());
        return schema;
    }

}
